package com.sjgs.gis.fs.store;

import com.google.common.collect.Lists;
import com.sjgs.common.utils.BytesUtil;
import com.sjgs.common.utils.DirUtils;
import com.sjgs.gis.common.RocksDBOptionsFactory;
import com.sjgs.gis.config.KVStoreProperties;
import com.sjgs.gis.errors.LocalFileIOException;
import com.sjgs.gis.errors.RocksIOException;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.rocksdb.*;

import java.io.Closeable;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import static com.sjgs.common.Constants.*;

/**
 * 瓦片kv存储
 * @author jxw
 */
public final class DefaultKVStore extends RocksStore implements Closeable {
    private static final Log logger = LogFactory.getLog(DefaultKVStore.class);
    private final ReadWriteLock readWriteLock = new ReentrantReadWriteLock();

    private List<ColumnFamilyOptions> cfOptionsList;
    private List<ColumnFamilyDescriptor> cfDescriptors;
    private Map<String, ColumnFamilyHandle> cfHandleMap;

    private RocksDB db;

    public DefaultKVStore(KVStoreProperties kvStoreProperties) {
        super(kvStoreProperties);
    }

    @Override
    public Boolean initStore() {
        cfOptionsList = new ArrayList<>();
        cfDescriptors = new ArrayList<>();
        ColumnFamilyOptions cfOptions = createColumnFamilyOptions();
        cfOptionsList.add(cfOptions);
        cfDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, cfOptions));
        cfDescriptors.add(new ColumnFamilyDescriptor(CF_LOCK.getBytes(), new ColumnFamilyOptions()));
        cfDescriptors.add(new ColumnFamilyDescriptor(CF_CLUSTER.getBytes(), new ColumnFamilyOptions()));
        cfDescriptors.add(new ColumnFamilyDescriptor(CF_BINLOG.getBytes(), new ColumnFamilyOptions()));
        cfDescriptors.add(new ColumnFamilyDescriptor(CF_SYNCLOG.getBytes(), new ColumnFamilyOptions()));
        return open();
    }

    private ColumnFamilyOptions createColumnFamilyOptions() {
        final BlockBasedTableConfig tConfig = RocksDBOptionsFactory.getRocksDBTableFormatConfig(DefaultKVStore.class);
        return RocksDBOptionsFactory.getRocksDBColumnFamilyOptions(DefaultKVStore.class) //
                .setTableFormatConfig(tConfig) //
                .setMergeOperator(new StringAppendOperator());
    }

    public Boolean open() {
        cfHandleMap = new HashMap<>();
        final List<ColumnFamilyHandle> cfHandles = Lists.newArrayList();
        KVStoreProperties kvStoreProperties = getKvStoreProperties();
        String dBpath = kvStoreProperties.getRootDir() + "/" + DEFAULT_DB;
        DirUtils.mkDir(dBpath);
        try {
            this.db = RocksDB.open(getOptions(), dBpath, this.cfDescriptors, cfHandles);
            this.cfHandleMap.put(new String(RocksDB.DEFAULT_COLUMN_FAMILY), cfHandles.get(0));
            this.cfHandleMap.put(CF_LOCK, cfHandles.get(1));
            this.cfHandleMap.put(CF_CLUSTER, cfHandles.get(2));
            this.cfHandleMap.put(CF_BINLOG, cfHandles.get(3));
            this.cfHandleMap.put(CF_SYNCLOG, cfHandles.get(4));
            return true;
        } catch (RocksDBException e) {
            throw new RocksIOException(e.getMessage());
        }
    }

    public ColumnFamilyHandle getDefaultColumnFamilyHandle() {
        return getColumnFamilyHandle(new String(RocksDB.DEFAULT_COLUMN_FAMILY));
    }

    public ColumnFamilyHandle getColumnFamilyHandle(String columnFamilyName) {
        return cfHandleMap.get(columnFamilyName);
    }

    @Override
    public void close() {
        Iterator it = cfHandleMap.keySet().iterator();
        while (it.hasNext()) {
            try (ColumnFamilyHandle columnFamilyHandle = cfHandleMap.get(it.next())) {
                if (columnFamilyHandle != null) {
                    columnFamilyHandle.close();
                }
            }
        }
        if (db != null) {
            db.close();
            db = null;
        }
    }

    public RocksDB getDb() {
        return db;
    }

     /**
     * 创建快照文件
     * @param sstFilePath
     * @param startKey
     * @param endKey
     * @return
     */
    CompletableFuture<Void> createSstFiles(String sstFilePath,final byte[] startKey,
                                           final byte[] endKey) {
        final ExecutorService executor = Executors.newSingleThreadExecutor();
        final CompletableFuture<Void> sstFuture = new CompletableFuture<>();
        final Snapshot snapshot = this.db.getSnapshot();
        if (getKvStoreProperties().isAsyncSnapshot()) {
            doCreateSstFile(sstFilePath,snapshot, startKey, endKey, sstFuture);
            return sstFuture;
        }
        // async snapshot
        executor.execute(() -> doCreateSstFile(sstFilePath,snapshot, startKey, endKey, sstFuture));
        return sstFuture;
    }

    public void doCreateSstFile(String sstFilePath,Snapshot snapshot, final byte[] startKey,
                                final byte[] endKey, CompletableFuture<Void> sstFuture) {
        final Lock readLock = this.readWriteLock.readLock();
        readLock.lock();
        try (final ReadOptions readOptions = new ReadOptions();
             final EnvOptions envOptions = new EnvOptions();
             final Options options = new Options().setMergeOperator(new StringAppendOperator())) {
            readOptions.setSnapshot(snapshot);
            try (final RocksIterator it = this.db.newIterator(getDefaultColumnFamilyHandle(), readOptions);
                 final SstFileWriter sstFileWriter = new SstFileWriter(envOptions, options)) {
                if (startKey == null) {
                    it.seekToFirst();
                } else {
                    it.seek(startKey);
                }
                sstFileWriter.open(sstFilePath);
                long count = 0;
                for (; ; ) {
                    if (!it.isValid()) {
                        break;
                    }
                    final byte[] key = it.key();
                    if (endKey != null && BytesUtil.compare(key, endKey) >= 0) {
                        break;
                    }
                    sstFileWriter.put(key, it.value());
                    ++count;
                    it.next();
                }
                if (count == 0) {
                    sstFileWriter.close();
                } else {
                    sstFileWriter.finish();
                }
                sstFuture.complete(null);
            } catch (final RocksDBException e) {
                throw new RocksIOException(e.getMessage());
            } finally {
                // Nothing to release, rocksDB never own the pointer for a snapshot.
                snapshot.close();
                // The pointer to the snapshot is released by the database instance.
                this.db.releaseSnapshot(snapshot);
            }
        }finally {
            readLock.unlock();
        }
    }

    /**
     * 通过快照加载数据
     * @param sstFile
     */
    void ingestSstFiles(String sstFile) {
        final Lock readLock = this.readWriteLock.readLock();
        readLock.lock();
        try (final IngestExternalFileOptions ingestOptions = new IngestExternalFileOptions()) {
            if (Files.size(Paths.get(sstFile)) == 0L) {
                return;
            }
            logger.info("Start ingest sst file :" + sstFile);
            this.db.ingestExternalFile(getDefaultColumnFamilyHandle(), Collections.singletonList(sstFile), ingestOptions);
        } catch (final RocksDBException e) {
            throw new LocalFileIOException("Fail to ingest sst file at path: " + sstFile);
        } catch (IOException e) {
            throw new LocalFileIOException("Fail to ingest sst file at path: " + sstFile);
        } finally {
            readLock.unlock();
        }
    }

    BackupInfo backupDB(String backupDBPath) throws IOException {
        final Lock writeLock = this.readWriteLock.writeLock();
        writeLock.lock();
        try (final BackupableDBOptions backupOpts = createBackupDBOptions(backupDBPath);
             final BackupEngine backupEngine = BackupEngine.open(this.getOptions().getEnv(), backupOpts)) {
            backupEngine.createNewBackup(this.db, true);
            final List<BackupInfo> backupInfoList = backupEngine.getBackupInfo();
            if (backupInfoList.isEmpty()) {
                logger.warn("Fail to backup at +" + backupDBPath);
                return null;
            }
            // chose the backupInfo who has max backupId
            return Collections.max(backupInfoList, Comparator.comparingInt(BackupInfo::backupId));
        } catch (RocksDBException e) {
            e.printStackTrace();
        } finally {
            writeLock.unlock();
        }
        return null;
    }

    void restoreBackup(final String backupDBPath, final BackupInfo rocksBackupInfo) {
        final Lock writeLock = this.readWriteLock.writeLock();
        writeLock.lock();
        close();
        try (final BackupableDBOptions backupOpts = createBackupDBOptions(backupDBPath);
             final BackupEngine backupEngine = BackupEngine.open(this.getOptions().getEnv(), backupOpts);
             final RestoreOptions restoreOpts = new RestoreOptions(false)) {
            final String dbPath = getKvStoreProperties().getRootDir() + "/" + DEFAULT_DB;
            backupEngine.restoreDbFromBackup(rocksBackupInfo.backupId(), dbPath, dbPath, restoreOpts);
            logger.info("Restored rocksDB from " + backupDBPath);
            // reopen the db
            open();
        } catch (final RocksDBException e) {
            throw new RocksIOException("Fail to restore from path  " + backupDBPath);
        } finally {
            writeLock.unlock();
        }
    }

    // Creates the backupable db options to control the behavior of
    // a backupable database.
    private static BackupableDBOptions createBackupDBOptions(final String backupDBPath) {
        return new BackupableDBOptions(backupDBPath) //
                .setSync(true) //
                .setShareTableFiles(false); // don't share data between backups
    }
}
